#include "protocol_detector.h"
#include "tds_protocol.h"
#include "mysql_protocol.h"
#include "oracle_protocol.h"
#include "postgreV3_protocol.h"
#include "dm_protocol.h"
#include "dm6_protocol.h"
#include "log.h"
#include <fstream>

//一个tcp会话识别包的个数范围
#define MIN_MATCH_COUNT 5
#define MAX_MATCH_COUNT 100

namespace secsmart_protocol
{
    ProtocolDetector::ProtocolDetector(std::string& targetFileName, uint32_t matchCount){
        LOG_DEBUG("targetFileName: {}; matchCount: {}",targetFileName, matchCount);
        if(MIN_MATCH_COUNT <= matchCount && matchCount <= MAX_MATCH_COUNT){
            maxMatchCount = matchCount;
        }else{
            maxMatchCount = MIN_MATCH_COUNT;
        }
        outFileName = targetFileName;
        init();
    };
    ProtocolDetector::~ProtocolDetector(){
       protocolLike.clear();
    }

    void ProtocolDetector::dissector(secsmart_session::session_sptr session, const uint8_t *buf, size_t size, bool is_forward){
        LOG_DEBUG("buf len: {}, detectState = {}",size,(int)detectState);
        if(detectState != DETECT_STATE::DETECT_STATE_INIT){
            return;
        }
        /*
        //当处理包的个数达到一定数量还无法识别，就不在识别
        if(parseCount > maxMatchCount * 2){
            detectState = DETECT_STATE::DETECT_STATE_UNKNOW;
            return;
        }
        */
        DETECT_RESULT result = DETECT_RESULT::DETECT_INIT;
        for(uint32_t i = 0; i < (uint32_t)protocolLike.size(); ++i){
            if(maxMatchCount < protocolLike[i]->getMatchCount() 
               || maxMatchCount < protocolLike[i]->getNoMatchCount()) {
                LOG_DEBUG("maxMatchCount={},MatchCount={},notMatchCount={}", 
                          maxMatchCount, protocolLike[i]->getMatchCount(), protocolLike[i]->getNoMatchCount())
                continue;
            }
            result = protocolLike[i]->parse(session, buf, size, is_forward);
            if(result == DETECT_RESULT::DETECT_OK){
                LOG_DEBUG("success: {}({}:{}) MatchCount:{} ",protocolLike[i]->getDatabaseName(), 
                                                 protocolLike[i]->getAssetIP(), 
                                                 protocolLike[i]->getAssetPort(),
                                                 protocolLike[i]->getMatchCount());
                if(maxMatchCount <= protocolLike[i]->getMatchCount() &&
                  protocolLike[i]->getDatabaseName().size() > 0){
                  detectState = DETECT_STATE::DETECT_STATE_SUCCESS;
                  writeFile(protocolLike[i]);
                }
                break;
            }
        }
        //++parseCount;
        
    };

    void ProtocolDetector::init(){
        protocolLike.emplace_back(std::make_shared<secsmart_protocol::MySQLProtocol>());
        protocolLike.emplace_back(std::make_shared<secsmart_protocol::TdsProtocol>());
        protocolLike.emplace_back(std::make_shared<secsmart_protocol::OracleProtocol>());
        protocolLike.emplace_back(std::make_shared<secsmart_protocol::PostgreProtocol>());
        //
        protocolLike.emplace_back(std::make_shared<secsmart_protocol::DMProtocol>());
        protocolLike.emplace_back(std::make_shared<secsmart_protocol::DM6Protocol>());
    };

    void ProtocolDetector::writeFile(std::shared_ptr<secsmart_protocol::BaseProtocol> protocol){
        if (outFileName.empty()){
            return;
        }
        std::string strOut = "{ \"DB_Type\":\"" + protocol->getDatabaseName() + "\", ";
        strOut += "\"DB_Version\":\"" + protocol->getVersion() + "\", ";
        strOut += "\"ip\":\""+ protocol->getAssetIP() + "\", ";
        strOut += "\"port\":\""+ std::to_string(protocol->getAssetPort()) + "\"";
        strOut += protocol->extendInfo();
        strOut += " }";

        std::ofstream ofs(outFileName, std::ios_base::app);
        ofs << strOut << std::endl;
    }

};